# 训练数字、小写字母、和大写字母
import tensorflow as tf
from tensorflow import keras
import gzip
import numpy as np

# 加载EMNIST数据集
def load_emnist(images_file, labels_file):
    with gzip.open(images_file, 'rb') as f:
        images = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1, 28, 28)
    with gzip.open(labels_file, 'rb') as f:
        labels = np.frombuffer(f.read(), np.uint8, offset=8)
    return images, labels

# 加载训练和测试数据
train_images, train_labels = load_emnist('emnist-gzip/emnist-byclass-train-images-idx3-ubyte.gz', 
                                         'emnist-gzip/emnist-byclass-train-labels-idx1-ubyte.gz')
test_images, test_labels = load_emnist('emnist-gzip/emnist-byclass-test-images-idx3-ubyte.gz', 
                                       'emnist-gzip/emnist-byclass-test-labels-idx1-ubyte.gz')

# 数据预处理
train_images = train_images.reshape((train_images.shape[0], 28, 28, 1)).astype('float32') / 255
test_images = test_images.reshape((test_images.shape[0], 28, 28, 1)).astype('float32') / 255

# 构建模型
model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(62, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 训练模型
model.fit(train_images, train_labels, epochs=10, validation_split=0.1)

# 保存模型
model.save('emnist_model.h5')